import logging
import mindspore as ms
from mindspore import nn
import numpy as np
#import torch.nn as nn
#from fastai.vision import *
from src.PositionAttention import *
from src.ms_Retsranformer import ResTranformer
from src.ms_resnet import resnet45
from src.ms_model import Model
from mindspore.ops import operations as P
#from modules.attention import *
#from modules.backbone import ResTranformer
#from modules.model import Model
#from modules.resnet import resnet45


class BaseVision(Model):
    def __init__(self, config):
        super().__init__(config)
        self.loss_weight = config.model_vision_loss_weight
        self.out_channels =  512
        self.cast = P.Cast()
        if config.model_vision_backbone == 'transformer':
            self.backbone = ResTranformer(config)
        else: self.backbone = resnet45()
        
        if config.model_vision_attention == 'position':
            mode = 'nearest'
            self.attention = PositionAttention(
                max_length=config.dataset_max_length + 1,  # additional stop token
                mode=mode,
            )
        

        self.cls = nn.Dense(self.out_channels, self.charset.num_classes, weight_init ='uniform',bias_init='uniform')

        #if config.model_vision_checkpoint is not None:
           # logging.info(f'Read vision model from {config.model_vision_checkpoint}.')
            #self.load(config.model_vision_checkpoint)

        

    def construct(self, images, *args):
        
        features = self.backbone(images)  # (N, E, H, W)
        
        attn_vecs, attn_scores = self.attention(features)  # (N, T, E), (N, T, H, W)#需要的

        logits = self.cls(attn_vecs) # (N, T, C)
        
        pt_lengths = self._get_length(logits)
        
        

        return {'feature': attn_vecs, 'logits': logits, 'pt_lengths': pt_lengths,
                'attn_scores': attn_scores, 'loss_weight':self.loss_weight, 'name': 'vision'}


        # features = ms.Tensor(np.ones((1,26,512)),ms.dtype.float32)
        # logits= ms.Tensor(np.ones((1,26,37)),ms.dtype.float32)
        # pt_lengths = ms.Tensor(5,ms.dtype.int16)
        # attn_scores = ms.Tensor(np.ones((1,26,8,32)),ms.dtype.float32)

        # return {'feature': features, 'logits': logits, 'pt_lengths': pt_lengths,
        #         'attn_scores': attn_scores, 'loss_weight':1.0, 'name': 'vision'}        
 
